import os

import cv2
import numpy as np

from thirdparties.wider_face_evaluation.evaluation import *
from modules.retinaface.detector import FaceDetector



# <Function: evaluate_widerface/>
def evaluate_widerface(
    face_detector, 
    val_images_path="./path_to_widerface_val_images", 
    val_annotations_path="./path_to_widerface_annotations_mats", 
    min_score_thresh=0.3,
    iou_thresh=0.5,
    show=False
    ):

    # pred = get_preds(pred)
    # norm_score(pred)

    facebox_list, event_list, file_list, hard_gt_list, medium_gt_list, easy_gt_list = get_gt_boxes(val_annotations_path)
    event_num = len(event_list)
    thresh_num = 1000
    settings = ['easy', 'medium', 'hard']
    setting_gts = [easy_gt_list, medium_gt_list, hard_gt_list]
    aps = []
    for setting_id in range(len(settings)):
        # different setting
        gt_list = setting_gts[setting_id]
        count_face = 0
        pr_curve = np.zeros((thresh_num, 2)).astype('float')
        # [hard, medium, easy]
        pbar = tqdm.tqdm(range(event_num))
        for i in pbar:
            pbar.set_description('Processing {}'.format(settings[setting_id]))
            event_name = str(event_list[i][0][0])
            img_list = file_list[i][0]
            # pred_list = pred[event_name]
            sub_gt_list = gt_list[i][0]
            # img_pr_info_list = np.zeros((len(img_list), thresh_num, 2))
            gt_bbx_list = facebox_list[i][0]

            for j in range(len(img_list)):                                
                # pred_info = pred_list[str(img_list[j][0][0])]
                
                image_path = os.path.join(val_images_path, event_name, str(img_list[j][0][0])+".jpg")
                image = cv2.imread(image_path, 1)
                boxes, scores, landmarks = detector.apply_image(image, min_score_thresh, iou_thresh)
                
                pred_info = np.concatenate((boxes, scores.reshape(-1, 1)), 1)
                pred_info[:,2] = pred_info[:,2]-pred_info[:,0]
                pred_info[:,3] = pred_info[:,3]-pred_info[:,1]
                pred_info = np.round(pred_info).astype(np.int32)
                
                gt_boxes = gt_bbx_list[j][0].astype('float')
                keep_index = sub_gt_list[j][0]
                count_face += len(keep_index)

                if show:
                    # draw the results
                    for box, score, landmark in zip(boxes, scores, landmarks):
                        cv2.rectangle( image, (box[0], box[1]), (box[2], box[3]), (0, 0, 255), 1 )
                        cv2.putText( image, "{:.4f}".format(score), (int(box[0]), int(box[1])+12), cv2.FONT_HERSHEY_DUPLEX, 0.5, (0, 0, 255) )
                        cv2.circle( image, (landmark[0][0], landmark[0][1]), 1, (0, 0, 255), 2)
                        cv2.circle( image, (landmark[1][0], landmark[1][1]), 1, (0, 0, 255), 2)
                        cv2.circle( image, (landmark[2][0], landmark[2][1]), 1, (0, 0, 255), 2)
                        cv2.circle( image, (landmark[3][0], landmark[3][1]), 1, (0, 0, 255), 2)
                        cv2.circle( image, (landmark[4][0], landmark[4][1]), 1, (0, 0, 255), 2)
                    # end-for
                                            
                    for gt_box in gt_boxes:
                        cv2.rectangle( image, (int(gt_box[0]), int(gt_box[1])), (int(gt_box[0]+gt_box[2]), int(gt_box[1]+gt_box[3])), (0, 255, 0), 1 )
                    # end-for
                    
                    cv2.imshow("image", image)
                    cv2.waitKey(0)
                # end-if                

                if len(gt_boxes) == 0 or len(pred_info) == 0:
                    continue
                ignore = np.zeros(gt_boxes.shape[0])
                if len(keep_index) != 0:
                    ignore[keep_index-1] = 1
                pred_recall, proposal_list = image_eval(pred_info, gt_boxes, ignore, iou_thresh)

                _img_pr_info = img_pr_info(thresh_num, pred_info, proposal_list, pred_recall)

                pr_curve += _img_pr_info
        pr_curve = dataset_pr_info(thresh_num, pr_curve, count_face)

        propose = pr_curve[:, 0]
        recall = pr_curve[:, 1]

        ap = voc_ap(recall, propose)
        aps.append(ap)

        cv2.destroyAllWindows()

    print("==================== Results ====================")
    print("Easy   Val AP: {}".format(aps[0]))
    print("Medium Val AP: {}".format(aps[1]))
    print("Hard   Val AP: {}".format(aps[2]))
    print("=================================================")

    pass
# <Function: /evaluate_widerface>

# <Function: evaluate_widerface/>
def evaluate_widerface_easy(
    face_detector, 
    val_images_path="./path_to_widerface_val_images", 
    val_annotations_path="./path_to_widerface_annotations_mats", 
    min_score_thresh=0.3,
    iou_thresh=0.5,
    show=False
    ):

    # pred = get_preds(pred)
    # norm_score(pred)

    facebox_list, event_list, file_list, hard_gt_list, medium_gt_list, easy_gt_list = get_gt_boxes(val_annotations_path)
    event_num = len(event_list)
    thresh_num = 1000
    settings = ['easy']
    setting_gts = [easy_gt_list]
    aps = []
    
    # get image count
    images_count_list = [int(0)] * len(settings)
    for setting_id in range(len(settings)):
        for i in range(event_num):
            img_list = file_list[i][0]
            for j in range(len(img_list)): 
                images_count_list[setting_id] += int(1)
    # end-for

    for setting_id in range(len(settings)):
        # different setting
        gt_list = setting_gts[setting_id]
        count_face = 0
        pr_curve = np.zeros((thresh_num, 2)).astype('float')
        # [hard, medium, easy]
        pbar = tqdm.tqdm(total=images_count_list[setting_id])
        pbar.set_description('Processing {}'.format(settings[setting_id]))
        for i in range(event_num):            
            event_name = str(event_list[i][0][0])
            img_list = file_list[i][0]
            sub_gt_list = gt_list[i][0]
            img_pr_info_list = np.zeros((len(img_list), thresh_num, 2))
            gt_bbx_list = facebox_list[i][0]
            for j in range(len(img_list)):
                image_path = os.path.join(val_images_path, event_name, str(img_list[j][0][0])+".jpg")
                image = cv2.imread(image_path, 1)
                boxes, scores, landmarks = detector.apply_image(image, min_score_thresh, iou_thresh)
                # 
                pred_info = np.concatenate((boxes, scores.reshape(-1, 1)), 1)
                pred_info[:,2] = pred_info[:,2]-pred_info[:,0]
                pred_info[:,3] = pred_info[:,3]-pred_info[:,1]
                pred_info = np.round(pred_info).astype(np.int32)
                # 
                gt_boxes = gt_bbx_list[j][0].astype('float')
                keep_index = sub_gt_list[j][0]
                count_face += len(keep_index)

                if show:
                    # draw the results
                    for box, score, landmark in zip(boxes, scores, landmarks):
                        cv2.rectangle( image, (box[0], box[1]), (box[2], box[3]), (0, 0, 255), 1 )
                        cv2.putText( image, "{:.4f}".format(score), (int(box[0]), int(box[1])+12), cv2.FONT_HERSHEY_DUPLEX, 0.5, (0, 0, 255) )
                        cv2.circle( image, (landmark[0][0], landmark[0][1]), 1, (0, 0, 255), 2)
                        cv2.circle( image, (landmark[1][0], landmark[1][1]), 1, (0, 0, 255), 2)
                        cv2.circle( image, (landmark[2][0], landmark[2][1]), 1, (0, 0, 255), 2)
                        cv2.circle( image, (landmark[3][0], landmark[3][1]), 1, (0, 0, 255), 2)
                        cv2.circle( image, (landmark[4][0], landmark[4][1]), 1, (0, 0, 255), 2)
                    # end-for
                                            
                    for gt_box in gt_boxes:
                        cv2.rectangle( image, (int(gt_box[0]), int(gt_box[1])), (int(gt_box[0]+gt_box[2]), int(gt_box[1]+gt_box[3])), (0, 255, 0), 1 )
                    # end-for
                    
                    cv2.imshow("image", image)
                    cv2.waitKey(0)
                # end-if

                pbar.update()

                if len(gt_boxes) == 0 or len(pred_info) == 0:
                    continue
                ignore = np.zeros(gt_boxes.shape[0])
                if len(keep_index) != 0:
                    ignore[keep_index-1] = 1
                pred_recall, proposal_list = image_eval(pred_info, gt_boxes, ignore, iou_thresh)

                _img_pr_info = img_pr_info(thresh_num, pred_info, proposal_list, pred_recall)

                pr_curve += _img_pr_info
                
                pass
            # end-for
            pass
        # end-for
        pbar.close()

        pr_curve = dataset_pr_info(thresh_num, pr_curve, count_face)

        propose = pr_curve[:, 0]
        recall = pr_curve[:, 1]

        ap = voc_ap(recall, propose)
        aps.append(ap)

        cv2.destroyAllWindows()

    print("=== Results =====================================")
    print(" Easy Val AP: {:.4f}".format(aps[0]))
    print("=================================================")

    pass
# <Function: /evaluate_widerface>

if __name__ == "__main__":
    detector = FaceDetector("./checkpoints/retinaface_2020_06_25_18_19_08/args.json")
    evaluate_widerface(
        face_detector=detector, 
        val_images_path="./data/wider_face/WIDER_val/images/", 
        val_annotations_path="./data/annotations/", 
        min_score_thresh=0.3,
        iou_thresh=0.5
        )
    pass

# python -m eval_widerface